# 导入ocr_nlp_processor 
from ocr_nlp_processor import ocr_nlp_processor
import re

# 提取数字
def getNumber(value):
    # 提取文本中的数字部分
    number_pattern = re.compile(r'\d+')
    result = number_pattern.search(value)
    
    if result:
        extracted_number = result.group()
        return f"'{extracted_number}"  # 添加'
    else:
        return value  # 如果没有找到数字，直接返回原始值

# 定义prompt和表头的映射
prompt_to_header = {
    "购买方名称是什么？": {"header": "购买方名称"},
    "购买方纳税人识别号是什么？": {"header": "购买方纳税人识别号", "process_value": getNumber},
    "购买方开户行及账号": {"header": "购买方开户行及账号"},
    "销售方名称是什么？": {"header": "销售方名称"},
    "销售方纳税人识别号是什么？": {"header": "销售方纳税人识别号", "process_value": getNumber},
    "销售方开户行及账号": {"header": "销售方开户行及账号"},
    "发票号码是多少?": {"header": "发票号码", "process_value": getNumber},
    "发票总金额是多少?": {"header": "发票金额", "process_value": getNumber},
    "开票日期是什么时间": {"header": "开票日期"},
    "收款人是谁": {"header": "收款人"},
    "开票人是谁": {"header": "开票人"},
}

# 输入文件目录
input_directory = './发票/所有文件'
# 输出文件目录
output_directory = './发票/已处理文件'
# 输出CSV文件路径
output_csv = './发票/识别结果.csv'

# 调用通用的OCR和NLP处理函数
ocr_nlp_processor(input_directory, output_directory, output_csv, prompt_to_header)
